import json
import time
from uuid import uuid4

from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.responses import JSONResponse
from contextvars import ContextVar
from conf.config import app_config
from core.exception.errors import BaseExceptionMixin
from core.logger import admin_logger, api_logger
from utils.common_helper import get_json_loads


req_ctx: ContextVar[Request] = ContextVar("request_ctx")


class AccessMiddleware(BaseHTTPMiddleware):

    async def dispatch(self, r: Request, call_next: RequestResponseEndpoint) -> Response:
        """记录接口请求日志"""
        # 设置请求上下文
        req_ctx.set(r)
        stime = time.time()
        rid = str(uuid4().hex)
        r.state.request_id = rid
        method = r.method.upper()
        log_info = f'{rid} | {r.client.host} | {method} {r.url}'
        req_log = ''

        as_details = app_config.LOG_API_DETAILS
        if as_details:
            headers = json.dumps(dict(r.headers), ensure_ascii=False)
            req_log = f'{log_info} | headers: {headers}'

        if as_details and method == 'POST' and 'application/json' in r.headers.get('Content-Type', '').lower():
            rbody = await r.body()
            rbody = get_json_loads(rbody, default=str(rbody))
            req_log += f' | req: {rbody}'

        logger = api_logger if r.url.path.startswith('/api') else admin_logger

        code = 500
        try:
            resp: Response = await call_next(r)
            code = resp.status_code
            if 'application/json' in r.headers.get('Content-Type', '').lower() and (r.url.path.startswith('/api') or r.url.path.startswith('/admin')):
                # 回退到迭代器方式
                response_body = b"".join([chunk async for chunk in resp.body_iterator])
                logger.info(f'{req_log or log_info} | code: {code} | response_body: {response_body.decode()}')
                resp = JSONResponse(
                    content=json.loads(response_body),
                    status_code=resp.status_code,
                    headers=dict(resp.headers)
                )
        except Exception as e:
            isinstance(e, BaseExceptionMixin) or logger.opt(exception=True).error(f'{log_info} | exception: {e}')
            raise e
        finally:
            elapsed = '{:.3f}s'.format(time.time() - stime)
            logger.info(f'{req_log or log_info} | code: {code} | elapsed: {elapsed}')

        resp.headers["X-Process-Time"] = elapsed
        resp.headers["X-Request-Id"] = rid
        return resp
